{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "# An overview of QuantTensor and QuantConv2d\n",
    "\n",
    "In this initial tutorial, we take a first look at `QuantTensor`, a basic data structure in Brevitas, and at `QuantConv2d`, a typical quantized layer. `QuantConv2d` is an instance of a `QuantWeightBiasInputOutputLayer` (typically imported as `QuantWBIOL`), meaning that it supports quantization of its weight, bias, input and output. Other instances of `QuantWBIOL` are `QuantLinear`, `QuantConv1d`, `QuantConvTranspose1d` and `QuantConvTranspose2d`, and they all follow the same principles.\n",
    "\n",
    "If we take a look at the `__init__` method of `QuantConv2d`, we notice a few things:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:36.342704Z",
     "iopub.status.busy": "2025-05-09T14:47:36.342005Z",
     "iopub.status.idle": "2025-05-09T14:47:49.604896Z",
     "shell.execute_reply": "2025-05-09T14:47:49.601319Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "    def __init__(\n",
       "            self,\n",
       "            in_channels: int,\n",
       "            out_channels: int,\n",
       "            kernel_size: Union[int, Tuple[int, int]],\n",
       "            stride: Union[int, Tuple[int, int]] = 1,\n",
       "            padding: Union[int, Tuple[int, int]] = 0,\n",
       "            dilation: Union[int, Tuple[int, int]] = 1,\n",
       "            groups: int = 1,\n",
       "            padding_mode: str = 'zeros',\n",
       "            bias: Optional[bool] = True,\n",
       "            weight_quant: Optional[WeightQuantType] = Int8WeightPerTensorFloat,\n",
       "            bias_quant: Optional[BiasQuantType] = None,\n",
       "            input_quant: Optional[ActQuantType] = None,\n",
       "            output_quant: Optional[ActQuantType] = None,\n",
       "            return_quant_tensor: bool = False,\n",
       "            device: Optional[torch.device] = None,\n",
       "            dtype: Optional[torch.dtype] = None,\n",
       "            **kwargs) -> None:\n",
       "        # avoid an init error in the super class by setting padding to 0\n",
       "        if padding_mode == 'zeros' and padding == 'same' and (stride > 1 if isinstance(\n",
       "                stride, int) else any(map(lambda x: x > 1, stride))):\n",
       "            padding = 0\n",
       "            is_same_padded_strided = True\n",
       "        else:\n",
       "            is_same_padded_strided = False\n",
       "        Conv2d.__init__(\n",
       "            self,\n",
       "            in_channels=in_channels,\n",
       "            out_channels=out_channels,\n",
       "            kernel_size=kernel_size,\n",
       "            stride=stride,\n",
       "            padding=padding,\n",
       "            padding_mode=padding_mode,\n",
       "            dilation=dilation,\n",
       "            groups=groups,\n",
       "            bias=bias,\n",
       "            device=device,\n",
       "            dtype=dtype)\n",
       "        QuantWBIOL.__init__(\n",
       "            self,\n",
       "            weight_quant=weight_quant,\n",
       "            bias_quant=bias_quant,\n",
       "            input_quant=input_quant,\n",
       "            output_quant=output_quant,\n",
       "            return_quant_tensor=return_quant_tensor,\n",
       "            **kwargs)\n",
       "        self.is_same_padded_strided = is_same_padded_strided\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import inspect\n",
    "from brevitas.nn import QuantConv2d\n",
    "from brevitas.nn import QuantIdentity\n",
    "from IPython.display import Markdown, display\n",
    "import torch\n",
    "\n",
    "# helpers\n",
    "def assert_with_message(condition):\n",
    "    assert condition\n",
    "    print(condition)\n",
    "\n",
    "def pretty_print_source(source):\n",
    "    display(Markdown('```python\\n' + source + '\\n```'))\n",
    "\n",
    "# set manual seed for the notebook\n",
    "torch.manual_seed(0)\n",
    "    \n",
    "source = inspect.getsource(QuantConv2d.__init__)  \n",
    "pretty_print_source(source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QuantConv2d` is an instance of both `Conv2d` and `QuantWBIOL`. Its initialization method exposes the usual arguments of a `Conv2d`, as well as: an extra flag to support *same padding*; *four* different arguments to set a quantizer for - respectively - *weight*, *bias*, *input*, and *output*; a `return_quant_tensor` boolean flag; the `**kwargs` placeholder to intercept additional arbitrary keyword arguments.  \n",
    "In this tutorial we will focus on how to set the four quantizer arguments and the return flags; arbitrary kwargs will be explained in a separate tutorial dedicated to defining and overriding quantizers.\n",
    "\n",
    "By default `weight_quant=Int8WeightPerTensorFloat`, while `bias_quant`, `input_quant` and `output_quant` are set to `None`. That means that by default weights are quantized to *8-bit signed integer with a per-tensor floating-point scale factor* (a very common type of quantization adopted by e.g. the ONNX standard opset), while quantization of bias, input, and output are disabled. We can easily verify all of this at runtime on an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:49.764711Z",
     "iopub.status.busy": "2025-05-09T14:47:49.763376Z",
     "iopub.status.idle": "2025-05-09T14:47:49.827557Z",
     "shell.execute_reply": "2025-05-09T14:47:49.825285Z"
    }
   },
   "outputs": [],
   "source": [
    "default_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:49.836377Z",
     "iopub.status.busy": "2025-05-09T14:47:49.835701Z",
     "iopub.status.idle": "2025-05-09T14:47:49.848377Z",
     "shell.execute_reply": "2025-05-09T14:47:49.846183Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Is weight quant enabled: True\n",
      "Is bias quant enabled: False\n",
      "Is input quant enabled: False\n",
      "Is output quant enabled: False\n"
     ]
    }
   ],
   "source": [
    "print(f'Is weight quant enabled: {default_quant_conv.weight_quant.is_quant_enabled}')\n",
    "print(f'Is bias quant enabled: {default_quant_conv.bias_quant.is_quant_enabled}')\n",
    "print(f'Is input quant enabled: {default_quant_conv.input_quant.is_quant_enabled}')\n",
    "print(f'Is output quant enabled: {default_quant_conv.output_quant.is_quant_enabled}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we now try to pass in a random floating-point tensor as input, as expected we get the output of the convolution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:49.857296Z",
     "iopub.status.busy": "2025-05-09T14:47:49.856602Z",
     "iopub.status.idle": "2025-05-09T14:47:50.139399Z",
     "shell.execute_reply": "2025-05-09T14:47:50.137028Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/_tensor.py:1255: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /opt/conda/conda-bld/pytorch_1670525541990/work/c10/core/TensorImpl.h:1758.)\n",
      "  return super(Tensor, self).rename(names)\n",
      "[W NNPACK.cpp:53] Could not initialize NNPACK! Reason: Unsupported hardware.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 0.2908, -0.1793, -0.9610],\n",
       "          [-0.6542, -0.3532,  0.6361],\n",
       "          [ 1.0290,  0.2730,  0.0969]],\n",
       "\n",
       "         [[-0.3479,  0.6030,  0.4900],\n",
       "          [ 0.1607,  0.3547, -0.4283],\n",
       "          [-0.6696,  0.0652,  0.7300]],\n",
       "\n",
       "         [[-0.0769, -0.2424,  0.1860],\n",
       "          [ 0.1740, -0.1182, -0.7017],\n",
       "          [ 0.0963,  0.2375, -0.9439]]]], grad_fn=<ConvolutionBackward0>)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "out = default_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case we are computing the convolution between an unquantized input tensor and quantized weights, so the output in general is unquantized."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A QuantConv2d with quantization disabled everywhere behaves like a standard `Conv2d`. Again can easily verify this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.151260Z",
     "iopub.status.busy": "2025-05-09T14:47:50.150107Z",
     "iopub.status.idle": "2025-05-09T14:47:50.201355Z",
     "shell.execute_reply": "2025-05-09T14:47:50.198237Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "from torch.nn import Conv2d\n",
    "\n",
    "torch.manual_seed(0)  # set a seed to make sure the random weight init is reproducible\n",
    "disabled_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, weight_quant=None)\n",
    "torch.manual_seed(0)  # reproduce the same random weight init as above\n",
    "float_conv = Conv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False)\n",
    "inp = torch.randn(1, 2, 5, 5)\n",
    "assert_with_message(torch.isclose(disabled_quant_conv(inp), float_conv(inp)).all().item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we have just seen, Brevitas allows users as much freedom as possible to experiment with quantization, meaning that computation between quantized and unquantized values is considered legal. This allows users to mix Brevitas layers with Pytorch layers with little restrictions.  \n",
    "To make this possible, quantized values are typically represented in *dequantized format*, meaning that - in the case of affine quantization implemented in Brevitas - zero-point and scale factor are applied to their integer values according to the formula **quant_value = (integer_value - zero_point) * scale**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QuantTensor\n",
    "\n",
    "We can directly observe the quantized weights by calling the weight quantizer on the layer's weights: `default_quant_conv.weight_quant(quant_conv.weight)`, which for shortness is already implemented as `default_quant_conv.quant_weight()` :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.212555Z",
     "iopub.status.busy": "2025-05-09T14:47:50.211864Z",
     "iopub.status.idle": "2025-05-09T14:47:50.239642Z",
     "shell.execute_reply": "2025-05-09T14:47:50.236802Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.0018,  0.1273, -0.1937],\n",
       "          [-0.1734, -0.0904,  0.0627],\n",
       "          [-0.0055,  0.1863, -0.0203]],\n",
       "\n",
       "         [[ 0.0627, -0.0720, -0.0461],\n",
       "          [-0.2251, -0.1568, -0.0978],\n",
       "          [ 0.0092,  0.0941,  0.1421]]],\n",
       "\n",
       "\n",
       "        [[[-0.1605, -0.1033,  0.0849],\n",
       "          [ 0.1956, -0.0480,  0.1771],\n",
       "          [-0.0387,  0.0258,  0.2140]],\n",
       "\n",
       "         [[-0.2196, -0.1476, -0.0590],\n",
       "          [-0.0923,  0.2030, -0.1531],\n",
       "          [-0.1089, -0.1642, -0.2214]]],\n",
       "\n",
       "\n",
       "        [[[-0.1384,  0.2030,  0.1052],\n",
       "          [ 0.1144,  0.0129, -0.1199],\n",
       "          [ 0.0406, -0.2196, -0.1697]],\n",
       "\n",
       "         [[-0.1218,  0.1494,  0.1384],\n",
       "          [-0.1052, -0.0092,  0.1513],\n",
       "          [ 0.2343,  0.0941,  0.0314]]]], grad_fn=<MulBackward0>), scale=tensor(0.0018, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "default_quant_conv.quant_weight()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the quantized weights are wrapped in a data structure implemented by Brevitas called `QuantTensor`. A `QuantTensor` is a way to represent an affine quantized tensor with all its metadata, meaning: the `value` of the quantized tensor in *dequantized* format, `scale`, `zero_point`, `bit_width`, whether the quantized value it's `signed` or not, and whether the tensor was generated in `training` mode. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, we have that the quantized value (in dequantized format) can be computer from its integer representation, together with zero-point and scale:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.253870Z",
     "iopub.status.busy": "2025-05-09T14:47:50.252982Z",
     "iopub.status.idle": "2025-05-09T14:47:50.283572Z",
     "shell.execute_reply": "2025-05-09T14:47:50.280681Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "int_weight = default_quant_conv.quant_weight().int()\n",
    "zero_point = default_quant_conv.weight_quant.zero_point()\n",
    "scale = default_quant_conv.weight_quant.scale()\n",
    "quant_weight_manually = (int_weight - zero_point) * scale\n",
    "\n",
    "assert_with_message(default_quant_conv.quant_weight().value.isclose(quant_weight_manually).all().item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A *valid* QuantTensor correctly populates all its fields with values `!= None` and respect the **affine quantization invariant**, i.e. `value / scale + zero_point` is (accounting for rounding errors) an *integer* that can be represented within the interval defined by the `bit_width` and `signed` fields of the `QuantTensor`. A *non-valid* one doesn't.\n",
    "We can observe that the quantized weights are indeed marked as valid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.294989Z",
     "iopub.status.busy": "2025-05-09T14:47:50.293979Z",
     "iopub.status.idle": "2025-05-09T14:47:50.309942Z",
     "shell.execute_reply": "2025-05-09T14:47:50.307679Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(default_quant_conv.quant_weight().is_valid)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling `is_valid` is relative expensive, so it should be using sparingly, but there are a few cases where a non-valid QuantTensor might be generated that is important to be aware of. Say we have two QuantTensor as output of the same quantized activation, and we want to sum them together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.321004Z",
     "iopub.status.busy": "2025-05-09T14:47:50.320270Z",
     "iopub.status.idle": "2025-05-09T14:47:50.358378Z",
     "shell.execute_reply": "2025-05-09T14:47:50.356059Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "tensor(0.0211, grad_fn=<AbsBinarySignGradFnBackward>)\n",
      "tensor(0.0162, grad_fn=<AbsBinarySignGradFnBackward>)\n"
     ]
    }
   ],
   "source": [
    "quant_act = QuantIdentity(return_quant_tensor=True)\n",
    "\n",
    "out_tensor_0 = quant_act(torch.randn(1,2,5,5))\n",
    "out_tensor_1 = quant_act(torch.randn(1,2,5,5))\n",
    "\n",
    "assert_with_message(out_tensor_0.is_valid)\n",
    "assert_with_message(out_tensor_1.is_valid)\n",
    "print(out_tensor_0.scale)\n",
    "print(out_tensor_1.scale)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both QuantTensor are valid but since the quantized activation is in training mode by default, their scale factors are going to be different. It is important to note that the behaviour is different at evaluation time, where the two scale factors will be the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.368684Z",
     "iopub.status.busy": "2025-05-09T14:47:50.367982Z",
     "iopub.status.idle": "2025-05-09T14:47:50.523192Z",
     "shell.execute_reply": "2025-05-09T14:47:50.519237Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IntQuantTensor(value=tensor([[[[-0.1106,  1.1945, -0.4972, -2.0968,  0.7175],\n",
      "          [-2.5901,  0.0588, -0.2014,  2.1486,  1.6435],\n",
      "          [ 0.9067, -2.5212,  2.2193,  0.2352, -0.8395],\n",
      "          [-0.8351,  0.6341, -0.5551,  0.1040, -3.3151],\n",
      "          [-0.8979, -0.7092,  3.8232,  1.0875,  0.3954]],\n",
      "\n",
      "         [[ 1.4363, -1.3973,  1.3249,  2.6914,  0.3660],\n",
      "          [ 1.5057,  1.8094,  0.5100, -1.6874,  1.9981],\n",
      "          [ 1.2472, -1.7813,  0.0334, -1.2880, -2.9333],\n",
      "          [ 0.0180, -1.4298, -2.9978,  0.5494, -1.4548],\n",
      "          [ 1.6738, -0.3177, -0.3721, -0.1650, -1.1871]]]],\n",
      "       grad_fn=<AddBackward0>), scale=0.018651068210601807, zero_point=0.0, bit_width=9.0, signed_t=True, training_t=True)\n"
     ]
    }
   ],
   "source": [
    "out_tensor = out_tensor_0 + out_tensor_1\n",
    "print(out_tensor)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we set `training` to `True` for both of them, we are allowed to sum them even if they have different scale factors. The output QuantTensor will have the correct `bit_width`, and a scale which is the average of the two original scale factors. This is done only at training time, in order to propagate gradient information, however the consequence is that the resulting QuantTensor is no longer valid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.534007Z",
     "iopub.status.busy": "2025-05-09T14:47:50.533247Z",
     "iopub.status.idle": "2025-05-09T14:47:50.550317Z",
     "shell.execute_reply": "2025-05-09T14:47:50.547377Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(not out_tensor.is_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QuantTensor` implements `__torch_function__` to handle being called from torch functional operators (e.g. ops under `torch.nn.functional`). Passing a QuantTensor to supported ops that are invariant to quantization, e.g. max-pooling, preserve the the validity of a QuantTensor. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.560831Z",
     "iopub.status.busy": "2025-05-09T14:47:50.559868Z",
     "iopub.status.idle": "2025-05-09T14:47:50.596944Z",
     "shell.execute_reply": "2025-05-09T14:47:50.594589Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[0.5191, 0.6402],\n",
       "          [2.1455, 0.5883]],\n",
       "\n",
       "         [[2.0417, 0.5883],\n",
       "          [1.2631, 0.3980]],\n",
       "\n",
       "         [[0.7959, 0.5191],\n",
       "          [0.8132, 1.3496]]]], grad_fn=<MaxPool2DWithIndicesBackward0>), scale=tensor(0.0173, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "quant_tensor = quant_identity(torch.randn(1, 3, 4, 4))\n",
    "torch.nn.functional.max_pool2d(quant_tensor, kernel_size=2, stride=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ops that are not invariant to quantization, a `QuantTensor` decays into a floating-point `torch.Tensor`. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.606781Z",
     "iopub.status.busy": "2025-05-09T14:47:50.606079Z",
     "iopub.status.idle": "2025-05-09T14:47:50.623164Z",
     "shell.execute_reply": "2025-05-09T14:47:50.620886Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 0.4770,  0.2212,  0.0691,  0.5650],\n",
       "          [-0.0346, -0.6618, -0.4635, -0.3482],\n",
       "          [ 0.9730, -0.7245, -0.5881, -0.5287],\n",
       "          [-0.0863,  0.8857,  0.5287, -0.4498]],\n",
       "\n",
       "         [[ 0.9669,  0.5650, -0.6211, -0.4498],\n",
       "          [-0.2376,  0.6103,  0.5287,  0.2700],\n",
       "          [-0.6808,  0.8519,  0.2700, -0.5531],\n",
       "          [-0.0173,  0.8264,  0.3782, -0.1881]],\n",
       "\n",
       "         [[-0.6211, -0.9764, -0.5993,  0.4770],\n",
       "          [ 0.5033,  0.6618, -0.1881, -0.6211],\n",
       "          [-0.8031,  0.1375,  0.5287,  0.8740],\n",
       "          [-0.6714,  0.6714, -0.5650,  0.8611]]]], grad_fn=<TanhBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.tanh(quant_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input Quantization\n",
    "\n",
    "We can obtain a valid output `QuantTensor` by making sure that both input and weight of `QuantConv2d` are quantized. To do so, we can set a quantizer for `input_quant`. In this example we pick a *signed 8-bit* quantizer with *per-tensor floating-point scale factor*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.632787Z",
     "iopub.status.busy": "2025-05-09T14:47:50.632045Z",
     "iopub.status.idle": "2025-05-09T14:47:50.697943Z",
     "shell.execute_reply": "2025-05-09T14:47:50.695138Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.3568, -0.1883,  0.3589],\n",
       "          [-0.4470,  0.1039, -0.3945],\n",
       "          [-0.4190,  0.3723,  0.8384]],\n",
       "\n",
       "         [[-0.0510,  0.5514, -0.2751],\n",
       "          [-0.5668,  0.5824,  0.2328],\n",
       "          [ 0.1316, -0.2518,  1.0418]],\n",
       "\n",
       "         [[ 0.2734,  0.7268, -0.0249],\n",
       "          [-0.1732,  0.5197,  1.1158],\n",
       "          [ 0.3771, -0.3810,  0.2008]]]], grad_fn=<ConvolutionBackward0>), scale=tensor([[[[3.1958e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor([0.]), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8ActPerTensorFloat\n",
    "\n",
    "input_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, \n",
    "    input_quant=Int8ActPerTensorFloat, return_quant_tensor=True)\n",
    "out_tensor = input_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.710406Z",
     "iopub.status.busy": "2025-05-09T14:47:50.708923Z",
     "iopub.status.idle": "2025-05-09T14:47:50.726441Z",
     "shell.execute_reply": "2025-05-09T14:47:50.723329Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(out_tensor.is_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens internally is that the input tensor passed to `input_quant_conv` is being quantized before being passed to the convolution operator. That means we are now computing a convolution between two quantized tensors, which mimplies that the output of the operation is also quantized. As expected then `out_tensor` is marked as valid. \n",
    "\n",
    "Another important thing to notice is how the `bit_width` field of `out_tensor` is relatively high at *21 bits*. In Brevitas, the assumption is always that the output bit-width of an operator reflects the worst-case size of the *accumulator* required by that operation. In other terms, given the *size* of the input and weight tensors and their *bit-widths*, 21 is the bit-width that would be required to represent the largest possible output value that could be generated. This makes sure that the affine quantization invariant is always respected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could have obtained a similar result by directly passing as input a QuantTensor. In this example we are directly defining a QuantTensor ourselves, but it could also be the output of a previous layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.739024Z",
     "iopub.status.busy": "2025-05-09T14:47:50.736526Z",
     "iopub.status.idle": "2025-05-09T14:47:50.774949Z",
     "shell.execute_reply": "2025-05-09T14:47:50.771768Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[ 7.2000e-03, -3.7000e-03,  7.7000e-03, -2.4000e-03, -8.9000e-03],\n",
       "          [-1.2000e-02, -8.1000e-03,  7.2000e-03, -1.1300e-02, -9.7000e-03],\n",
       "          [-1.0000e-03,  1.0100e-02,  3.8000e-03, -1.1900e-02,  6.9000e-03],\n",
       "          [ 8.3000e-03,  1.0000e-04, -6.9000e-03,  3.9000e-03, -5.4000e-03],\n",
       "          [ 1.1300e-02, -6.0000e-03,  9.7000e-03,  0.0000e+00,  1.0900e-02]],\n",
       "\n",
       "         [[-1.0900e-02,  1.1400e-02, -6.4000e-03,  9.2000e-03,  7.1000e-03],\n",
       "          [-6.0000e-04,  9.2000e-03, -8.5000e-03,  5.0000e-03,  6.5000e-03],\n",
       "          [-8.3000e-03, -1.2000e-03,  7.4000e-03,  9.2000e-03, -6.0000e-04],\n",
       "          [-2.1000e-03,  9.5000e-03,  3.0000e-04, -2.9000e-03, -6.5000e-03],\n",
       "          [-1.1800e-02, -4.8000e-03,  5.4000e-03, -2.5000e-03,  9.0000e-04]]]]), scale=tensor(1.0000e-04), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant_tensor import IntQuantTensor\n",
    "\n",
    "scale = 0.0001\n",
    "bit_width = 8\n",
    "zero_point = 0.\n",
    "int_value = torch.randint(low=- 2 ** (bit_width - 1), high=2 ** (bit_width - 1) - 1, size=(1, 2, 5, 5))\n",
    "quant_value = (int_value - zero_point) * scale\n",
    "quant_tensor_input = IntQuantTensor(\n",
    "    quant_value, \n",
    "    scale=torch.tensor(scale), \n",
    "    zero_point=torch.tensor(zero_point), \n",
    "    bit_width=torch.tensor(float(bit_width)),\n",
    "    signed=True,\n",
    "    training=True)\n",
    "quant_tensor_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.787168Z",
     "iopub.status.busy": "2025-05-09T14:47:50.785661Z",
     "iopub.status.idle": "2025-05-09T14:47:50.804242Z",
     "shell.execute_reply": "2025-05-09T14:47:50.800920Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(quant_tensor_input.is_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: how we are explicitly forcing `value`, `scale`, `zero_point` and `bit_width` to be floating-point `torch.Tensor`, as this is expected by Brevitas but it's currently not enforced automatically at initialization time.\n",
    "\n",
    "If we now pass in `quant_tensor_input` to `return_quant_conv`, we will see that indeed the output is a valid 21-bit `QuantTensor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.816035Z",
     "iopub.status.busy": "2025-05-09T14:47:50.814764Z",
     "iopub.status.idle": "2025-05-09T14:47:50.868306Z",
     "shell.execute_reply": "2025-05-09T14:47:50.865095Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.0019,  0.0049, -0.0012],\n",
       "          [-0.0012,  0.0050, -0.0074],\n",
       "          [-0.0023, -0.0035, -0.0033]],\n",
       "\n",
       "         [[-0.0031,  0.0028,  0.0116],\n",
       "          [ 0.0079,  0.0046,  0.0022],\n",
       "          [ 0.0021, -0.0004,  0.0011]],\n",
       "\n",
       "         [[-0.0045, -0.0010,  0.0002],\n",
       "          [-0.0044,  0.0027,  0.0025],\n",
       "          [-0.0009,  0.0040, -0.0044]]]], grad_fn=<ConvolutionBackward0>), scale=tensor([[[[1.8307e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor([0.]), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "return_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, return_quant_tensor=True)\n",
    "out_tensor = return_quant_conv(quant_tensor_input)\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.880509Z",
     "iopub.status.busy": "2025-05-09T14:47:50.879500Z",
     "iopub.status.idle": "2025-05-09T14:47:50.893788Z",
     "shell.execute_reply": "2025-05-09T14:47:50.890925Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(out_tensor.is_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also pass in an input `QuantTensor` to a layer that has `input_quant` enabled. In that case, the input gets re-quantized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.903845Z",
     "iopub.status.busy": "2025-05-09T14:47:50.902962Z",
     "iopub.status.idle": "2025-05-09T14:47:50.929161Z",
     "shell.execute_reply": "2025-05-09T14:47:50.927067Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.0073,  0.0040, -0.0011],\n",
       "          [-0.0033,  0.0078, -0.0028],\n",
       "          [ 0.0005, -0.0025, -0.0008]],\n",
       "\n",
       "         [[ 0.0021, -0.0021,  0.0035],\n",
       "          [ 0.0012, -0.0016, -0.0023],\n",
       "          [-0.0010, -0.0015,  0.0040]],\n",
       "\n",
       "         [[-0.0010,  0.0047,  0.0025],\n",
       "          [-0.0014,  0.0021, -0.0039],\n",
       "          [ 0.0036, -0.0003,  0.0026]]]], grad_fn=<ConvolutionBackward0>), scale=tensor([[[[1.7393e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor([0.]), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_quant_conv(quant_tensor_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Output Quantization\n",
    "\n",
    "Let's now look at would have happened if we instead enabled output quantization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:50.941833Z",
     "iopub.status.busy": "2025-05-09T14:47:50.941080Z",
     "iopub.status.idle": "2025-05-09T14:47:50.995902Z",
     "shell.execute_reply": "2025-05-09T14:47:50.993032Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.2117, -0.4811,  0.0385],\n",
       "          [-0.5100, -0.2502, -0.2213],\n",
       "          [-0.5773,  0.0192, -0.5485]],\n",
       "\n",
       "         [[ 0.1347,  0.8179, -1.2316],\n",
       "          [-0.6062,  0.4426, -0.3849],\n",
       "          [ 0.1732, -0.5100, -0.1251]],\n",
       "\n",
       "         [[ 1.0873,  0.2406, -0.2887],\n",
       "          [-0.4330, -0.4907, -0.2021],\n",
       "          [ 0.6447,  0.4811,  0.1347]]]], grad_fn=<MulBackward0>), scale=tensor(0.0096, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8ActPerTensorFloat\n",
    "\n",
    "output_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, \n",
    "    output_quant=Int8ActPerTensorFloat, return_quant_tensor=True)\n",
    "out_tensor = output_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:51.006408Z",
     "iopub.status.busy": "2025-05-09T14:47:51.005506Z",
     "iopub.status.idle": "2025-05-09T14:47:51.020016Z",
     "shell.execute_reply": "2025-05-09T14:47:51.017295Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(out_tensor.is_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see again that the output is a valid `QuantTensor`. However, what happened internally is quite different from before.  \n",
    "Previously, we computed the convolution between two quantized tensors, and got a quantized tensor as output.  \n",
    "In this case instead, we compute the convolution between a quantized and an unquantized tensor, we take its unquantized output and we quantize it.  \n",
    "The difference is obvious once we look at the output `bit_width`. In the previous case, we had that the `bit_width` reflected the size of the output accumulator. In this case instead, we have `bit_width=tensor(8.)`, which is what we expected since `output_quant` had been set to an *Int8* quantizer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bias Quantization\n",
    "\n",
    "There is an important scenario where the various options we just saw make a practical difference, and it's quantization of *bias*. In many contexts, such as in the ONNX standard opset and in FINN, bias is assumed to be quantized with scale factor equal to *input scale * weight scale*, which means that we need a valid quantized input somehow. A predefined bias quantizer that reflects that assumption is `brevitas.quant.scaled_int.Int8Bias`. If we simply tried to set it to a `QuantConv2d` without any sort of input quantization, we would get an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:51.030976Z",
     "iopub.status.busy": "2025-05-09T14:47:51.030102Z",
     "iopub.status.idle": "2025-05-09T14:47:53.184432Z",
     "shell.execute_reply": "2025-05-09T14:47:53.182221Z"
    },
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "QuantLayer is not correctly configured",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[23], line 6\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mbrevitas\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mquant\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mscaled_int\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Int8Bias\n\u001b[1;32m      3\u001b[0m bias_quant_conv \u001b[38;5;241m=\u001b[39m QuantConv2d(\n\u001b[1;32m      4\u001b[0m     in_channels\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m, out_channels\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m3\u001b[39m, kernel_size\u001b[38;5;241m=\u001b[39m(\u001b[38;5;241m3\u001b[39m,\u001b[38;5;241m3\u001b[39m), bias\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m      5\u001b[0m     bias_quant\u001b[38;5;241m=\u001b[39mInt8Bias, return_quant_tensor\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m----> 6\u001b[0m \u001b[43mbias_quant_conv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrandn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py:1194\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1190\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1191\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1192\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1195\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1196\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_conv.py:194\u001b[0m, in \u001b[0;36mQuantConv2d.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    193\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Union[Tensor, QuantTensor]) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Union[Tensor, QuantTensor]:\n\u001b[0;32m--> 194\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_layer.py:152\u001b[0m, in \u001b[0;36mQuantWeightBiasInputOutputLayer.forward_impl\u001b[0;34m(self, inp)\u001b[0m\n\u001b[1;32m    148\u001b[0m compute_output_quant_tensor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28misinstance\u001b[39m(quant_input, QuantTensor) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\n\u001b[1;32m    149\u001b[0m     quant_weight, QuantTensor)\n\u001b[1;32m    150\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (compute_output_quant_tensor \u001b[38;5;129;01mor\u001b[39;00m\n\u001b[1;32m    151\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_quant\u001b[38;5;241m.\u001b[39mis_quant_enabled) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreturn_quant_tensor:\n\u001b[0;32m--> 152\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mQuantLayer is not correctly configured\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    154\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbias \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    155\u001b[0m     quant_bias \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbias_quant(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbias, quant_input, quant_weight)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: QuantLayer is not correctly configured"
     ]
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8Bias\n",
    "\n",
    "bias_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    bias_quant=Int8Bias, return_quant_tensor=True)\n",
    "bias_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can solve the issue by passing in a valid `QuantTensor`, e.g. the `quant_tensor_input`  we defined above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.193089Z",
     "iopub.status.busy": "2025-05-09T14:47:53.192699Z",
     "iopub.status.idle": "2025-05-09T14:47:53.212033Z",
     "shell.execute_reply": "2025-05-09T14:47:53.210109Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-2.4238e-03, -5.6598e-03,  5.1882e-03],\n",
       "          [-6.5582e-03,  8.9274e-03,  4.9640e-04],\n",
       "          [ 9.6283e-03, -1.7466e-03, -4.8311e-03]],\n",
       "\n",
       "         [[ 2.9322e-03, -3.1358e-03, -6.2727e-04],\n",
       "          [ 2.8723e-06, -3.7981e-03,  1.0973e-02],\n",
       "          [-4.1031e-03,  6.5909e-03, -4.2369e-03]],\n",
       "\n",
       "         [[ 4.1967e-03, -7.0733e-03,  1.6456e-03],\n",
       "          [ 1.8197e-03, -3.1683e-03,  4.8200e-03],\n",
       "          [-3.2585e-04,  3.1055e-03,  1.9703e-03]]]],\n",
       "       grad_fn=<ConvolutionBackward0>), scale=tensor([[[[1.7953e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor([0.]), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bias_quant_conv(quant_tensor_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or by enabling input quantization and then passing in a float a `torch.Tensor` or a `QuantTensor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.223026Z",
     "iopub.status.busy": "2025-05-09T14:47:53.222646Z",
     "iopub.status.idle": "2025-05-09T14:47:53.264483Z",
     "shell.execute_reply": "2025-05-09T14:47:53.262449Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.2816, -0.5271, -0.1748],\n",
       "          [-0.4247, -0.1575,  0.0681],\n",
       "          [ 0.6528, -0.5346, -0.0657]],\n",
       "\n",
       "         [[ 0.2993, -0.3383,  0.3035],\n",
       "          [-0.4595, -0.6796, -0.9720],\n",
       "          [-0.1948, -0.5169, -0.2175]],\n",
       "\n",
       "         [[ 0.5586,  0.0665, -0.5807],\n",
       "          [ 0.5565,  0.1780, -0.0555],\n",
       "          [-0.1080,  0.0791, -0.2262]]]], grad_fn=<ConvolutionBackward0>), scale=tensor([[[[4.2009e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor([0.]), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_bias_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    input_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias, return_quant_tensor=True)\n",
    "input_bias_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.276894Z",
     "iopub.status.busy": "2025-05-09T14:47:53.274767Z",
     "iopub.status.idle": "2025-05-09T14:47:53.298962Z",
     "shell.execute_reply": "2025-05-09T14:47:53.296871Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.0058,  0.0030,  0.0030],\n",
       "          [-0.0013, -0.0002,  0.0043],\n",
       "          [-0.0061,  0.0033, -0.0001]],\n",
       "\n",
       "         [[ 0.0013, -0.0008, -0.0015],\n",
       "          [ 0.0011,  0.0012, -0.0012],\n",
       "          [-0.0013, -0.0020,  0.0002]],\n",
       "\n",
       "         [[-0.0061,  0.0053, -0.0004],\n",
       "          [ 0.0028,  0.0031, -0.0037],\n",
       "          [ 0.0027, -0.0048, -0.0044]]]], grad_fn=<ConvolutionBackward0>), scale=tensor([[[[1.7370e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor([0.]), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_bias_quant_conv(quant_tensor_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the output `bit_width=tensor(22.)`. This is because, in the worst-case, summing a *21-bit* integer (the size of the accumulator before bias is added) and an *8-bit* integer (the size of quantized bias) gives a *22-bit* integer.\n",
    "\n",
    "Let's try now to enable output quantization instead of input quantization. That wouldn't have solved the problem with bias quantization, as output quantization is performed after bias is added:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.310569Z",
     "iopub.status.busy": "2025-05-09T14:47:53.309670Z",
     "iopub.status.idle": "2025-05-09T14:47:53.518628Z",
     "shell.execute_reply": "2025-05-09T14:47:53.516652Z"
    },
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/parameter_quant.py:190: UserWarning: No quant bias cache found, set cache_inference_quant_bias=True and run an inference pass first\n",
      "  warn(\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "Input scale required",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[27], line 4\u001b[0m\n\u001b[1;32m      1\u001b[0m output_bias_quant_conv \u001b[38;5;241m=\u001b[39m QuantConv2d(\n\u001b[1;32m      2\u001b[0m     in_channels\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m, out_channels\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m3\u001b[39m, kernel_size\u001b[38;5;241m=\u001b[39m(\u001b[38;5;241m3\u001b[39m,\u001b[38;5;241m3\u001b[39m), bias\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m      3\u001b[0m     output_quant\u001b[38;5;241m=\u001b[39mInt8ActPerTensorFloat, bias_quant\u001b[38;5;241m=\u001b[39mInt8Bias, return_quant_tensor\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m----> 4\u001b[0m \u001b[43moutput_bias_quant_conv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrandn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py:1194\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1190\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1191\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1192\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1195\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1196\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_conv.py:194\u001b[0m, in \u001b[0;36mQuantConv2d.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    193\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Union[Tensor, QuantTensor]) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Union[Tensor, QuantTensor]:\n\u001b[0;32m--> 194\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_layer.py:155\u001b[0m, in \u001b[0;36mQuantWeightBiasInputOutputLayer.forward_impl\u001b[0;34m(self, inp)\u001b[0m\n\u001b[1;32m    152\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mQuantLayer is not correctly configured\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    154\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbias \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 155\u001b[0m     quant_bias \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias_quant\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mquant_input\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mquant_weight\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    156\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    157\u001b[0m     quant_bias \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py:1194\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1190\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1191\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1192\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1195\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1196\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/parameter_quant.py:364\u001b[0m, in \u001b[0;36mBiasQuantProxyFromInjector.forward\u001b[0;34m(self, x, input, weight)\u001b[0m\n\u001b[1;32m    362\u001b[0m     input_scale \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mscale()\n\u001b[1;32m    363\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m input_scale \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 364\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInput scale required\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    365\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrequires_input_scale \u001b[38;5;129;01mand\u001b[39;00m input_scale \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mis_quant_enabled:\n\u001b[1;32m    366\u001b[0m     input_scale \u001b[38;5;241m=\u001b[39m input_scale\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Input scale required"
     ]
    }
   ],
   "source": [
    "output_bias_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    output_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias, return_quant_tensor=True)\n",
    "output_bias_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not all scenarios require bias quantization to depend on the scale factor of the input. In those cases, biases can be quantized the same way weights are quantized, and have their own scale factor. In Brevitas, a predefined quantizer that reflects this other scenario is `Int8BiasPerTensorFloatInternalScaling`. In this case then a valid quantized input is not required:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.527707Z",
     "iopub.status.busy": "2025-05-09T14:47:53.527316Z",
     "iopub.status.idle": "2025-05-09T14:47:53.555501Z",
     "shell.execute_reply": "2025-05-09T14:47:53.553106Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.4360, -0.2674, -0.4194],\n",
       "          [-0.2412, -0.6360, -0.6838],\n",
       "          [-0.5227, -0.0199, -0.1445]],\n",
       "\n",
       "         [[-0.3524,  0.8025,  0.2844],\n",
       "          [ 0.9945, -0.4782,  0.8064],\n",
       "          [ 0.5732,  0.1249,  0.3110]],\n",
       "\n",
       "         [[ 0.3223,  0.2530,  0.2753],\n",
       "          [ 0.5764, -0.2533, -0.0181],\n",
       "          [-0.4147,  0.2049, -0.9944]]]], grad_fn=<ConvolutionBackward0>)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8BiasPerTensorFloatInternalScaling\n",
    "\n",
    "bias_internal_scale_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    bias_quant=Int8BiasPerTensorFloatInternalScaling, return_quant_tensor=False)\n",
    "bias_internal_scale_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are a couple of situations to be aware of concerning bias quantization that can lead to changes in the output `zero_point`.\n",
    "\n",
    "Let's consider the scenario where we compute the convolution between a quantized input tensor and quantized weights. In the first case, we then add an *unquantized* bias on top of the output. In the second one, we add a bias quantized with its own scale factor, e.g. with the `Int8BiasPerTensorFloatInternalScaling` quantizer. In both cases, in order to make sure the output `QuantTensor` is valid (i.e. the affine quantization invariant is respected), the output `zero_point` becomes non-zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.568572Z",
     "iopub.status.busy": "2025-05-09T14:47:53.567742Z",
     "iopub.status.idle": "2025-05-09T14:47:53.612433Z",
     "shell.execute_reply": "2025-05-09T14:47:53.609734Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.6912,  0.0086,  0.1628],\n",
       "          [-0.4786, -0.8073,  0.5224],\n",
       "          [ 0.4157,  0.4686,  0.2560]],\n",
       "\n",
       "         [[ 0.3170, -0.5486, -0.5216],\n",
       "          [ 0.1832,  1.0217, -0.3637],\n",
       "          [-0.1115,  0.6974, -0.0452]],\n",
       "\n",
       "         [[-0.6168, -0.5241, -0.6593],\n",
       "          [ 0.6408,  0.2674,  0.4537],\n",
       "          [-0.3744, -0.7771, -0.2848]]]], grad_fn=<ConvolutionBackward0>), scale=tensor([[[[3.0094e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor([[[[  339.3406]],\n",
       "\n",
       "         [[-4597.1797]],\n",
       "\n",
       "         [[-3452.3713]]]], grad_fn=<DivBackward0>), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unquant_bias_input_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    input_quant=Int8ActPerTensorFloat, return_quant_tensor=True)\n",
    "out_tensor = unquant_bias_input_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.624581Z",
     "iopub.status.busy": "2025-05-09T14:47:53.623748Z",
     "iopub.status.idle": "2025-05-09T14:47:53.637044Z",
     "shell.execute_reply": "2025-05-09T14:47:53.635022Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(out_tensor.is_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, an important point about `QuantTensor`. With the exception of learned bit-width (which will be the subject of a separate tutorial) and some of the bias quantization scenarios we have just seen, usually returing a `QuantTensor` is not necessary and can create extra complexity. This is why currently `return_quant_tensor` defaults to `False`. We can easily see it in an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:47:53.649111Z",
     "iopub.status.busy": "2025-05-09T14:47:53.648714Z",
     "iopub.status.idle": "2025-05-09T14:47:53.682037Z",
     "shell.execute_reply": "2025-05-09T14:47:53.680350Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2327,  0.9267,  0.6294],\n",
       "          [ 0.0901,  0.1027, -0.0727],\n",
       "          [-0.5614,  0.6182,  0.5394]],\n",
       "\n",
       "         [[ 0.4179, -0.5184, -0.2016],\n",
       "          [ 0.1390, -0.3925, -0.6171],\n",
       "          [ 0.4782,  0.0814,  0.6124]],\n",
       "\n",
       "         [[ 0.2896, -0.3779,  0.9408],\n",
       "          [-0.1334,  0.6186,  0.2167],\n",
       "          [-0.5926,  0.3690, -0.0284]]]], grad_fn=<ConvolutionBackward0>)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bias_input_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    input_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias)\n",
    "bias_input_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Altough not obvious, the output is actually implicitly quantized."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "98d39f4fefefd06bad749c114549705efc94e31e1348b8fb7b25329049b354fa"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
